# 人工智能NLP-Agent数字人项目-04-基金数据问答任务工单V1.1-20250214
import csv
import re
import pandas as pd
import utils.configFinRAG as configFinRAG
from utils.instances import LLM, TOKENIZER
from utils import prompts
from langchain_core.prompts import ChatPromptTemplate


# 封装 Jaccard 相似度计算函数
def jaccard_similarity(list1, list2):
    intersection = len(set(list1) & set(list2))
    union = len(set(list1)) + len(set(list2))
    return intersection / union if union != 0 else 0


def generate_answer(user_question, query_result, language_model, example_questions, example_infos, example_answers,
                    example_token_lists, num_examples=5):
    # 过滤 8 位数字的正则表达式
    number_pattern = r'\d{8}'
    processed_question = user_question
    # 提取数字
    numbers = re.findall(number_pattern, processed_question)
    # 将数字替换为空格
    for number in numbers:
        processed_question = processed_question.replace(number, ' ')
    # 对处理后的问题进行 token 化
    question_tokens = TOKENIZER(processed_question)['input_ids']
    # 计算与已有问题的相似度
    similarities = [jaccard_similarity(question_tokens, token_list) for token_list in example_token_lists]
    # 获取相似度最高的 num_examples 个问题的索引
    top_indices = sorted(range(len(similarities)), key=lambda i: similarities[i], reverse=True)[:num_examples]
    # 防止提示语过长
    prompt_length = 0
    selected_indices = []
    for index in top_indices:
        prompt_length += len(example_questions[index]) + len(example_answers[index])
        if prompt_length > 2000:
            break
        selected_indices.append(index)
    # 组装 prompt
    prompt_template = ChatPromptTemplate.from_template(prompts.ANSWER_TEMPLATE)
    example_prompts = '\n'.join([
        f"问题：{example_questions[index]}\n资料：{example_infos[index]}\n答案：{example_answers[index]}"
        for index in selected_indices
    ])
    chain = prompt_template | language_model
    response = chain.invoke({"examples": example_prompts, "FA": query_result, "question": user_question})
    return response.content


def load_example_data():
    # 读取问题和 FA 模板
    example_data = pd.read_csv(configFinRAG.sql_examples_path, delimiter=",", header=0)
    questions = example_data['问题'].tolist()
    infos = example_data['资料'].tolist()
    answers = example_data['FA'].tolist()
    token_lists = [TOKENIZER(question)['input_ids'] for question in questions]
    return questions, infos, answers, token_lists


if __name__ == '__main__':
    try:
        # 第一步：读取问题和 FA 模板，使用 tokenizer 进行 token 化
        example_questions, example_infos, example_answers, example_token_lists = load_example_data()
        # 第二步：读取问题执行结果文件
        result_data = pd.read_csv(configFinRAG.question_sql_check_path, delimiter=",", header=0)
        # 打开答案文件准备写入
        with open(configFinRAG.answer_path, 'w', newline='', encoding='utf-8-sig') as answer_file:
            csv_writer = csv.writer(answer_file)
            csv_writer.writerow(['问题id', '问题', '资料', 'FA'])
            # 第三步：循环问题，使用 Jaccard 进行相似度计算
            for _, row in result_data.iterrows():
                if row['flag'] == 1:
                    answer = generate_answer(
                        row['问题'], row['执行结果'], LLM,
                        example_questions, example_infos, example_answers,
                        example_token_lists
                    )
                    csv_writer.writerow([
                        str(row['问题id']),
                        str(row['问题']),
                        str(row['执行结果']),
                        answer
                    ])
    except Exception as e:
        print(f"程序运行出错: {e}")
